from math import ceil
from numba import cuda
import numpy as np


def get_default_device():
    """
    获取默认的GPU设备
    :return: 当前设备
    """
    return cuda.get_current_device()


def set_default_device(device_id=0):
    """
    设置默认的GPU设备，必须在调用之前设置
    :param device_id: 设备id
    :return:当前设备
    """
    return cuda.select_device(device_id)


def list_devices():
    """
    获取当前可用的设备列表
    :return: 当前可用的设备列表
    """
    return cuda.list_devices()


def get_grid_1d(x):
    assert isinstance(x, int) and x > 0

    block_dim_x = min(x, 1024)
    block_dim = (block_dim_x,)

    grid_dim_x = ceil(x / block_dim_x)
    grid_dim = (grid_dim_x,)

    assert block_dim_x * grid_dim_x <= 2147483647
    return grid_dim, block_dim


def get_grid_2d(x, y):
    assert isinstance(x, int) and x > 0
    assert isinstance(y, int) and y > 0

    block_dim_x = x if x * y <= 1024 else 32
    block_dim_y = y if x * y <= 1024 else 32
    block_dim = (block_dim_x, block_dim_y)

    grid_dim_x = ceil(x / block_dim_x)
    grid_dim_y = ceil(y / block_dim_y)
    grid_dim = (grid_dim_x, grid_dim_y)

    assert block_dim_x * grid_dim_x <= 2147483647
    assert block_dim_y * grid_dim_y <= 65535

    return grid_dim, block_dim


def get_grid_3d(x, y, z):
    assert isinstance(x, int) and x > 0
    assert isinstance(y, int) and y > 0
    assert isinstance(z, int) and z > 0

    block_dim_x = x if x * y * z <= 1024 else 16
    block_dim_y = y if x * y * z <= 1024 else 16
    block_dim_z = z if x * y * z <= 1024 else 4
    block_dim = (block_dim_x, block_dim_y, block_dim_z)

    grid_dim_x = ceil(x / block_dim_x)
    grid_dim_y = ceil(y / block_dim_y)
    grid_dim_z = ceil(z / block_dim_z)
    grid_dim = (grid_dim_x, grid_dim_y, grid_dim_z)

    return grid_dim, block_dim


def assume_ndarray(*ls: np.ndarray):
    if len(ls) == 0:
        return
    elif len(ls) == 1:
        x = ls[0]
        return x if isinstance(x, np.ndarray) else np.array(x)
    else:
        return [x if isinstance(x, np.ndarray) else np.array(x) for x in ls]


def flatten(x: np.ndarray):
    return x.reshape((-1,))
